package cn.foolishbird.crow.structure.skiplist;

import java.util.Random;

class SkipListNode {
    int key;
    int value;
    SkipListNode next;
    SkipListNode down;

    public SkipListNode(int key, int value) {
        this.key = key;
        this.value = value;
    }
}

public class SkipList {
    private static final double PROBABILITY = 0.5; // 跳跃概率
    private SkipListNode head;
    private int size;
    private int height;
    private Random random;

    public SkipList() {
        head = new SkipListNode(Integer.MIN_VALUE, Integer.MIN_VALUE);
        head.next = new SkipListNode(Integer.MAX_VALUE, Integer.MAX_VALUE);
        head.down = null;
        size = 0;
        height = 0;
        random = new Random();
    }

    public void put(int key, int value) {
        SkipListNode cur = head;

        // 记录需要更新的节点
        SkipListNode[] update = new SkipListNode[height + 1];
        for (int i = height; i >= 0; i--) {
            while (cur.next != null && cur.next.key < key) {
                cur = cur.next;
            }
            update[i] = cur; // 更新节点
            cur = cur.down; // 向下一层移动
        }

        cur = cur.next;

        // 如果 key 已存在，更新其值并返回
        if (cur != null && cur.key == key) {
            cur.value = value;
            return;
        }

        // 插入新节点
        SkipListNode newNode = new SkipListNode(key, value);
        insertNode(update, newNode);

        // 根据概率决定是否添加新层
        while (random.nextDouble() < PROBABILITY) {
            if (height >= update.length) {
                addNewLayer();
            }

            while (height >= 0 && update[height] == null) {
                height--;
            }

            if (height >= 0) {
                update[height] = update[height].down;
            }

            insertNode(update, newNode);
        }

        size++;
    }

    public Integer get(int key) {
        SkipListNode cur = head;

        for (int i = height; i >= 0; i--) {
            while (cur.next != null && cur.next.key < key) {
                cur = cur.next;
            }
            if (cur.next != null && cur.next.key == key) {
                return cur.next.value;
            }
            cur = cur.down;
        }

        return null;
    }

    public void remove(int key) {
        SkipListNode cur = head;

        // 记录需要更新的节点
        SkipListNode[] update = new SkipListNode[height + 1];
        for (int i = height; i >= 0; i--) {
            while (cur.next != null && cur.next.key < key) {
                cur = cur.next;
            }
            update[i] = cur; // 更新节点
            cur = cur.down; // 向下一层移动
        }

        cur = cur.next;

        // 如果 key 不存在，直接返回
        if (cur == null || cur.key != key) {
            return;
        }

        // 逐层删除节点
        while (cur != null) {
            update[cur.key] = cur.next;
            cur = cur.down;
        }

        // 移除空层
        while (height > 0 && head.next == null) {
            head = head.down;
            height--;
        }

        size--;
    }

    private void insertNode(SkipListNode[] update, SkipListNode newNode) {
        for (int i = 0; i < update.length; i++) {
            newNode.next = update[i].next;
            update[i].next = newNode;
            newNode.down = null;
            newNode = newNode.down;
        }
    }

    private void addNewLayer() {
        SkipListNode newHead = new SkipListNode(Integer.MIN_VALUE, Integer.MIN_VALUE);
        SkipListNode newTail = new SkipListNode(Integer.MAX_VALUE, Integer.MAX_VALUE);
        newHead.next = newTail;
        newHead.down = head;
        newTail.down = head.next;
        head = newHead;
        height++;
    }

    public void printSkipList() {
        SkipListNode cur = head;

        while (cur != null) {
            SkipListNode node = cur;
            while (node != null) {
                System.out.print("(" + node.key + ", " + node.value + ") -> ");
                node = node.next;
            }
            System.out.println("null");
            cur = cur.down;
        }
    }

    public static void main(String[]args) {
        SkipList skipList = new SkipList();

        skipList.put(1, 10);
        skipList.put(2, 20);
        skipList.put(3, 30);
        skipList.put(4, 40);
        skipList.put(5, 50);

        skipList.printSkipList();

        System.out.println("Value for key 3: " + skipList.get(3));
        System.out.println("Value for key 6: " + skipList.get(6));

        skipList.remove(2);
        skipList.remove(4);

        skipList.printSkipList();
    }
}